{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "538e0c69-9a0d-4207-b591-dfc2d6e834b4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b830ed27-c57d-473a-85dc-4ee7df50b498",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d6fc5be7-b424-4f3c-bb49-1e70f8f0b95b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 查看当前tensor占用的显存\n",
    "torch.cuda.memory_allocated()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4ca5c77f-28c1-4927-9413-8da59421fe3e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9932111872"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 查看进程占用的总共的显存\n",
    "torch.cuda.memory_reserved()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "90aa7f92-744d-4b02-baed-5ceb07ced489",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 释放掉未使用的缓存\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cfb92a5c-2f42-4b7c-945b-78c8bbdc932b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "CUDA extension not installed.\n",
      "CUDA extension not installed.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0b34f1585aa94ce19d6bd1ec324205ad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "torch.manual_seed(1234)\n",
    "# Note: The default behavior now has injection attack prevention off.\n",
    "vl_tokenizer = AutoTokenizer.from_pretrained(\"/root/share/zpt/projects/automoment/models/Qwen-VL-Chat-Int4\", trust_remote_code=True)\n",
    "# use cuda device\n",
    "vl_model = AutoModelForCausalLM.from_pretrained(\"/root/share/zpt/projects/automoment/models/Qwen-VL-Chat-Int4\", device_map=\"cuda\", trust_remote_code=True).eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9e373f17-ce24-4cc6-a828-6db670582402",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QWenTokenizer(name_or_path='/root/share/zpt/projects/automoment/models/Qwen-VL-Chat-Int4', vocab_size=151860, model_max_length=8192, is_fast=False, padding_side='right', truncation_side='right', special_tokens={}, clean_up_tokenization_spaces=True)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vl_tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c3b24a5-5fba-4b60-a8dc-54ace2712ba8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "923b4c39-b4d1-4874-8c92-0d5583f228f5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<ref>坦克</ref><box>(577,228),(672,315)</box><box>(577,384),(670,487)</box><box>(426,378),(519,491)</box><box>(425,250),(508,341)</box><box>(367,243),(437,332)</box>\n",
      "<ref>坦克</ref><box>(433,335),(474,420)</box><box>(466,337),(512,419)</box>\n",
      "<ref>坦克</ref><box>(436,491),(583,631)</box><box>(328,425),(421,546)</box><box>(728,689),(888,796)</box><box>(287,607),(377,711)</box><box>(423,671),(532,766)</box><box>(399,57),(478,131)</box><box>(50,813),(188,905)</box>\n",
      "<ref>坦克</ref><box>(438,553),(534,662)</box><box>(422,505),(534,591)</box>\n",
      "<ref>坦克</ref><box>(327,246),(437,345)</box><box>(436,379),(538,490)</box><box>(582,378),(658,475)</box><box>(524,236),(627,333)</box><box>(304,645),(490,804)</box>\n",
      "<ref>坦克</ref><box>(329,481),(403,602)</box><box>(524,512),(608,629)</box>\n",
      "<ref>坦克</ref><box>(397,375),(438,496)</box><box>(503,375),(545,493)</box><box>(569,382),(660,493)</box><box>(441,375),(483,493)</box>\n",
      "<ref>坦克</ref><box>(520,396),(608,492)</box><box>(316,393),(409,480)</box>\n",
      "<ref>坦克</ref><box>(584,234),(687,334)</box><box>(660,230),(740,326)</box><box>(579,403),(666,508)</box><box>(358,234),(457,329)</box><box>(409,360),(484,462)</box>\n"
     ]
    }
   ],
   "source": [
    "def check_tank():\n",
    "    for file in os.listdir('imgs/tank'):\n",
    "        if not file.endswith('.jpg'):\n",
    "            continue\n",
    "        query = vl_tokenizer.from_list_format([\n",
    "            {'image': 'imgs/tank/' + file},\n",
    "            {'text': '输出所有\"坦克\"的检测框,至少五个'},\n",
    "        ])\n",
    "        response, history = vl_model.chat(vl_tokenizer, query=query, history=None)\n",
    "\n",
    "        print(response)\n",
    "        image = vl_tokenizer.draw_bbox_on_latest_picture(response, history)\n",
    "        if image:\n",
    "            image.save('imgs/tank/res/' + file)\n",
    "        else:\n",
    "            print(\"no box\")\n",
    "        # break\n",
    "check_tank()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc3236ce-3a9b-4f45-82af-b5f546566335",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vision",
   "language": "python",
   "name": "vision"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
